import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import numpy as np

from plotly.offline import init_notebook_mode
init_notebook_mode()

df = pd.read_csv('../data/capture_fisheries.csv', skiprows=4)
df = df.dropna(how='all')

df['Country Name'] = df['Country Name'].str.replace('"', '')
years = [str(year) for year in range(1960, 2025)]
id_vars = ["Country Name", "Country Code", "Indicator Name", "Indicator Code"]

df_melted = df.melt(id_vars=id_vars, value_vars=years, 
                   var_name="Year", value_name="Production")
df_melted['Year'] = pd.to_numeric(df_melted['Year'])
df_melted['Production'] = pd.to_numeric(df_melted['Production'], errors='coerce')
df_melted = df_melted.dropna(subset=['Production'])

q95 = df_melted['Production'].quantile(0.95)
focused_data = df_melted[df_melted['Production'] <= q95]
zmax = focused_data['Production'].max()
zmin = focused_data['Production'].min()

enhanced_scale = [
    [0.0, '#E1F5FE'], 
    [0.1, '#B3E5FC'],  
    [0.2, '#81D4FA'], 
    [0.3, '#4FC3F7'], 
    [0.4, '#29B6F6'], 
    [0.5, '#039BE5'], 
    [0.6, '#0288D1'], 
    [0.7, '#0277BD'], 
    [0.8, '#01579B'],  
    [0.9, '#0D47A1'],  
    [1.0, '#1d1c47']  
]


frames = []
for year in sorted(df_melted['Year'].unique()):
    yearly_data = df_melted[df_melted['Year'] == year]
    yearly_focused = yearly_data[yearly_data['Production'] <= yearly_data['Production'].quantile(0.95)]
    year_zmax = yearly_focused['Production'].max()
    year_zmin = yearly_focused['Production'].min()
    
    frames.append(go.Frame(
        data=[go.Choropleth(
            locations=yearly_data['Country Code'],
            z=yearly_data['Production'],
            zmin=year_zmin,
            zmax=year_zmax,
            colorscale=enhanced_scale,
            customdata=yearly_data[['Country Name']],
            hovertemplate="<b>%{customdata[0]}</b><br>" +
                         "Year: %{frame.name}<br>" +
                         "Production: %{z:,} metric tons<extra></extra>",
            marker_line_color='darkgray',
            marker_line_width=0.5
        )],
        name=str(year)
    ))

fig_slider = go.Figure(
    data=[go.Choropleth(
        locations=df_melted['Country Code'],
        z=df_melted['Production'],
        colorscale=enhanced_scale,
        customdata=df_melted[['Country Name']],
        hovertemplate="<b>%{customdata[0]}</b><br>" +
                     "Year: %{frame.name}<br>" +
                     "Production: %{z:,} metric tons<extra></extra>",
        marker_line_color='darkgray',
        marker_line_width=0.5
    )],
    layout=go.Layout(
        width=750,
        height=500,
        geo=dict(
            domain={'x': [0, 0.9], 'y': [0, 1]},
            showframe=True,
            showcoastlines=True,
            landcolor='#f0f0f0',
            lakecolor='#d0e0f0',
            bgcolor='white',
            framecolor='black',
            framewidth=1
        ),
        margin={'t': 30, 'r': 20, 'b': 200, 'l': 30},
        title={
        'text': '<b>Annual Fisheries Production per Country</b>',
        'x': 0.42,
        'xanchor': 'center',
        'font': {
            'color': '#0a2463',
        }
        },
        sliders=[{
            "active": 0,
            "steps": [{
                "args": [[f.name], {"frame": {"duration": 300, "redraw": True},
                                  "mode": "immediate"}],
                "label": f.name,
                "method": "animate"
            } for f in frames],
            "x": 0.15,
            "len": 0.7,
            "currentvalue": {
                "prefix": "<b>Year: </b>",
                "font": {"size": 14},
                "xanchor": "center"
            }
        }],
        updatemenus=[{
            "buttons": [
                {
                    "args": [None, {"frame": {"duration": 400, "redraw": True},
                                  "fromcurrent": True}],
                    "label": "Play",
                    "method": "animate"
                },
                {
                    "args": [[None], {"frame": {"duration": 0, "redraw": True},
                                    "mode": "immediate"}],
                    "label": "Pause",
                    "method": "animate"
                }
            ],
            "type": "buttons",
            "x": 0.1,
            "y": -0.05,
            "bgcolor": "rgba(255,255,255,0.8)",
            "bordercolor": "#444"
        }]
    ),
    frames=frames,
)

fig_slider.add_annotation(x=0, y=-0.60,
                   xref="paper", yref="paper",
                   showarrow=False,
                   align='left',
                   xanchor='left', yanchor='bottom',
                   text="Annual fisheries production per country since 1960.<br>" + \
                        'Use slider to see different years or press play to see time progression.<br>' + \
                        'Color scale automatically focuses on 95% of data for each year. <br>')

fig_slider.update_annotations(
    font=dict(color='#0a2463')
)

fig_slider.show()